package CompressTool

import (
	"archive/zip"
	"fmt"
	"gitee.com/dadjin/my-commonly-used-go-func/MyTool/FileTool"
	OtherToolAlexmullinsZip "gitee.com/dadjin/my-commonly-used-go-func/OtherTool/alexmullins/zip"
	"io"
	"os"
	"path/filepath"
)

// EncryptZip 加密压缩文件
// src: 源文件路径
// desc: 压缩文件路径
// password: <PASSWORD>
func EncryptZip(src, desc, password string) error {
	zipfile, err := os.Create(desc)
	if err != nil {
		return err
	}
	defer zipfile.Close()

	archive := OtherToolAlexmullinsZip.NewWriter(zipfile)
	defer archive.Close()

	if FileTool.IsDir(src) {
		filepath.Walk(src, func(path string, info os.FileInfo, err error) error {

			if err != nil {
				return err
			}

			header, err := OtherToolAlexmullinsZip.FileInfoHeader(info)

			if err != nil {
				return err
			}
			header.Name, err = filepath.Rel(src, path)
			if err != nil {
				return err
			}
			if info.IsDir() {
				header.Name += "/"
			} else {
				header.Method = zip.Deflate
			}

			// 设置密码
			return addEncryptFileToZip(path, header, archive, password)

		})

	} else {
		fileInfo, err := os.Stat(src)
		if err != nil {
			fmt.Println("无法获取文件信息:", err)
			return err
		}

		header, err := OtherToolAlexmullinsZip.FileInfoHeader(fileInfo)
		header.Method = zip.Deflate
		if err != nil {
			return err
		}
		return addEncryptFileToZip(src, header, archive, password)
	}
	return err

}

func addEncryptFileToZip(path string, header *OtherToolAlexmullinsZip.FileHeader, archive *OtherToolAlexmullinsZip.Writer, pwd string) error {

	header.SetPassword(pwd)
	writer, err := archive.CreateHeader(header)
	if err != nil {
		return err
	}
	if FileTool.IsFile(path) {
		file, err := os.Open(path)
		if err != nil {
			return err
		}
		defer file.Close()
		_, err = io.Copy(writer, file)
	}
	return err
}
